#include "meshImportExport.hpp"
#include "rectMeshFunctions.hpp"
#include "structMeshFunctions.hpp"
#include "macros.hpp"
#include "meshTypes.hpp"
#include "util.hpp"
#include "topology.hpp"


namespace hms
{


/**
 * @brief Returns a 3xn array, whose first row holds the indices of all boundary
 * vertices, and the second and third row holds the indices of the boundary
 * edges at those vertices.
 * 
 * @param edgeVertices 
 * @param nInnerEdges 
 * @return Array3Xi 
 */
Array3Xi edgesAtBoundaryVertices(
	const Array2Xi& edgeVertices,
	Index nInnerEdges
){
	Index nBoundaryEdges { edgeVertices.cols() - nInnerEdges };

	/* this array will hold vertex indices in the first row and edge indices
	 * in the second row. Every boundary vertex has two boundary edges attached
	 * to it, so there will be two times as many columns as boundary vertices.
	 * Later, it is to be sorted and both edges per vertex go into one column.
	 * */
	std::vector<Array2i> edgesAtVertices;
	edgesAtVertices.reserve( 2*nBoundaryEdges );

	for ( int i{0}; i<nBoundaryEdges; ++i ){
		Index edge { nInnerEdges+i };
		edgesAtVertices.emplace_back( edgeVertices(0, edge), edge );
		edgesAtVertices.emplace_back( edgeVertices(1, edge), edge );
	}

	Eigen::Map<Array2Xi> eigenMap { edgesAtVertices[0].data(), 2, 2*nBoundaryEdges };

	std::sort(
		edgesAtVertices.begin(),
		edgesAtVertices.end(),
		[](const auto& a, const auto& b){
			return a[0] != b[0] ? a[0] < b[0] : a[1] < b[1];
		}
	);

	/* in the resulting array, every vertex index appears twice right after one
	 * another in the first row, with the associated edge indices in the second
	 * row. With column-major storage, this means that rearranging this array
	 * into four rows, the first and third row contain the vertex indices and
	 * have the same entries. So, to get rid of the duplicate vertex indices,
	 * the third row can be dropped. */
	Eigen::Map<Array4Xi> map { edgesAtVertices[0].data(), 4, nBoundaryEdges };

	Array3Xi result (3, nBoundaryEdges);

	result.topRows<2>() = map.topRows<2>();
	result.bottomRows<1>() = map.bottomRows<1>();

	return result;
}

Array3Xi edgesAtBoundaryVertices( const Mesh& mesh ){
	return edgesAtBoundaryVertices( mesh.edges().vertices(), mesh.nInnerEdges() );
}


ArrayXi verticesToBoundarySegment(
	const Array2Xi& verticesAtEdge,
	const Array3Xi& edgesAtVertex,
	Index firstVertex,
	Index secondVertex,
	Index lastVertex
){
	std::vector<Index> edges;

	/* The first row of edgesAtVertex holds the boundary vertices in ascending
	 * order. The second and third row hold the boundary edges connected to that
	 * vertex. This lambda calls std::lower_bound to find the column belonging
	 * to a vertex in an efficient fashion. */
	auto findCol = [&](Index vertex) -> Array3Xi::ConstColXpr {
		return edgesAtVertex.col( std::lower_bound(
			edgesAtVertex.row(0).begin(),
			edgesAtVertex.row(0).end(),
			vertex
		) - edgesAtVertex.row(0).begin() );
	};

	Array3Xi::ConstColXpr edgesAtFirstVertex { findCol(firstVertex) };

	/* store the first segment edge by checking which edge at the first vertex
	 * connects to the second vertex */
	edges.push_back(
		sameEntries(
			verticesAtEdge.col(edgesAtFirstVertex[1]), 
			Array2i{firstVertex, secondVertex}
		) ? edgesAtFirstVertex[1] : edgesAtFirstVertex[2]
	);

	/* make sure that the input vertices have actually been found. Otherwise, 
	 * they may be invalid input */
	if ( !sameEntries(
		verticesAtEdge.col( edges.back() ),
		Array2i{ firstVertex, secondVertex }
	)){
		throw InputError(
			"Boundary segment: first and/or second vertex could not be found! "
			"First vertex index: " + std::to_string(firstVertex) + 
			", second vertex index: " + std::to_string(secondVertex)
		);
	}

	Index currentVertex { secondVertex };
	while (true){
		/* now we loop to find all edges until we arrive at the last vertex 
		 * or have fully looped around*/
		if (currentVertex == lastVertex || currentVertex == firstVertex)
			break;
		
		Array3Xi::ConstColXpr currentEdges = findCol(currentVertex);
		/* next edge is the one connected to the current vertex that hasn't
		 * already just been stored */
		Index nextEdge {
			currentEdges[1] != edges.back() ?
			currentEdges[1] :
			currentEdges[2]
		};
		edges.push_back(nextEdge);

		/* the next vertex is the one on the other end of the next edge */
		currentVertex = otherVertex( verticesAtEdge, nextEdge, currentVertex);
	}

	return toEigen<ArrayXi>(edges);
}

ArrayXi verticesToBoundarySegment(
	const Array2Xi& edgeVertices,
	const Array3Xi& edgesAtVertex,
	const Array3i& verts
){
	return verticesToBoundarySegment(
		edgeVertices, edgesAtVertex, verts[0], verts[1], verts[2]
	);
}




ArrayXi boundarySegment(
	BoundaryLocation loc,
	Index vBeg,
	Index vEnd,
	const Ref<const Array2i>& meshSize
){
	assert( vBeg < vEnd && "Use reverse version instead!" );
	Index firstEdge, nEdges;
	// using enum BoundaryLocation;
	switch(loc){
		case BoundaryLocation::bottom:
		case BoundaryLocation::top:
			firstEdge = vBeg - vBeg/(meshSize.x()+1);
			nEdges = vEnd - vBeg;
			return ArrayXi::LinSpaced(
				nEdges,
				firstEdge,
				firstEdge + nEdges - 1
			);
		case BoundaryLocation::left:
		case BoundaryLocation::right:
			firstEdge = nHorizontalEdges(meshSize) + vBeg;
			nEdges = (vEnd - vBeg) / (meshSize.x()+1);
			return ArrayXi::LinSpaced(
				nEdges,
				firstEdge,
				firstEdge + (nEdges - 1) * (meshSize.x()+1)
			);
	}
	return {};
}


ArrayXi boundarySegmentReverse(
	BoundaryLocation loc,
	Index vBeg,
	Index vEnd,
	const Ref<const Array2i>& meshSize
){
	assert( vBeg > vEnd && "Use non-reverse version instead!" );
	Index firstEdge, nEdges;
	// using enum BoundaryLocation;
	switch(loc){
		case BoundaryLocation::bottom:
		case BoundaryLocation::top:
			firstEdge = vBeg - (vBeg-1)/(meshSize.x()+1) - 1;
			nEdges = vBeg - vEnd;
			return ArrayXi::LinSpaced(
				nEdges,
				firstEdge,
				firstEdge - nEdges + 1
			);
		case BoundaryLocation::left:
		case BoundaryLocation::right:
			firstEdge = nHorizontalEdges(meshSize) + vBeg - (meshSize.x()+1);
			nEdges = (vBeg - vEnd) / (meshSize.x()+1);
			return ArrayXi::LinSpaced(
				nEdges,
				firstEdge,
				firstEdge - (nEdges - 1) * (meshSize.x()+1)
			);
	}
	return {};
}


ArrayXi verticesToBoundarySegment(
	const Ref<const Array2i>& meshSize,
	const Array3i& verts
){
	Index
		nx { meshSize.x() },
		ny { meshSize.y() };

	enum corner {
		bottomLeft = 0,
		bottomRight,
		topLeft,
		topRight
	};

	Array4i corners {
		0,              // bottom left
		nx,             // bottom right
		(nx+1)*ny,      // top left
		(nx+1)*(ny+1)-1 // top right
	};

	if ( (verts < 0).any() || (verts > corners[topRight]).any() )
		throw InputError( "Vertices out of bounds!" );

	/* find out which side of the mesh the first edge is at */
	BoundaryLocation locBeg;
	bool notAtBoundary {false};
	if ( abs(verts[1]-verts[0]) == 1 ){
		if ( (verts.head(2) <= corners[bottomRight]).all() )
			locBeg = BoundaryLocation::bottom;
		else if ( (verts.head(2) >= corners[topLeft]).all() )
			locBeg = BoundaryLocation::top;
		else
			notAtBoundary = true;
	} else if ( abs(verts[1]-verts[0]) == nx+1 ){
		if ( verts[0] % (nx+1) == 0 && verts[1] % (nx+1) == 0 )
			locBeg = BoundaryLocation::left;
		else if ( verts[0] % (nx+1) == nx && verts[1] % (nx+1) == nx )
			locBeg = BoundaryLocation::right;
		else
			notAtBoundary = true;
	} else {
		throw InputError(
			"Incorrect difference between first and second vertex!"
		);
	}

	bool
		endAtBottom {false},
		endAtTop {false},
		endAtLeft {false},
		endAtRight {false};

	if ( verts[2] <= corners[bottomRight] )
		endAtBottom = true;
	if ( verts[2] >= corners[topLeft] )
		endAtTop = true;
	if ( verts[2] % (nx+1) == 0 )
		endAtLeft = true;
	if ( verts[2] % (nx+1) == nx )
		endAtRight = true;
	
	if ( !( endAtBottom || endAtTop || endAtLeft || endAtRight ) )
		notAtBoundary = true;

	if (notAtBoundary)
		throw InputError( "Vertices not at boundary!");
	
	bool clockwise {
		(locBeg == BoundaryLocation::bottom || locBeg == BoundaryLocation::right) ?
			verts[1] < verts[0] : verts[1] > verts[0]
	};

	Index
		vBeg { verts[0] },
		vEnd { verts[2] };
	
	auto sameSide = [&](){
		if (clockwise){
			switch(locBeg){
				case BoundaryLocation::bottom:
					return verts[2] < vBeg;
				case BoundaryLocation::top:
					return verts[2] > vBeg;
				case BoundaryLocation::left:
					return verts[2] > vBeg && endAtLeft;
				case BoundaryLocation::right:
					return verts[2] < vBeg && endAtRight;
			}
		} else {
			switch(locBeg){
				case BoundaryLocation::bottom:
					return verts[2] > vBeg && endAtBottom;
				case BoundaryLocation::top:
					return verts[2] < vBeg && endAtTop;
				case BoundaryLocation::left:
					return verts[2] < vBeg && endAtLeft;
				case BoundaryLocation::right:
					return verts[2] > vBeg && endAtRight;
			}
		}
		throw std::runtime_error("Unexpected branch");
	};

	auto progressLoc = [&](){
		switch(locBeg){
			case BoundaryLocation::bottom:
				locBeg = clockwise ? BoundaryLocation::left : BoundaryLocation::right;
				break;
			case BoundaryLocation::left:
				locBeg = clockwise ? BoundaryLocation::top : BoundaryLocation::bottom;
				break;
			case BoundaryLocation::top:
				locBeg = clockwise ? BoundaryLocation::right : BoundaryLocation::left;
				break;
			case BoundaryLocation::right:
				locBeg = clockwise ? BoundaryLocation::bottom : BoundaryLocation::top;
		}
	};

	auto nextCorner = [&](){
		switch(locBeg){
			case BoundaryLocation::bottom:
				return clockwise ? corners[bottomLeft]  : corners[bottomRight];
			case BoundaryLocation::left:
				return clockwise ? corners[topLeft]     : corners[bottomLeft];
			case BoundaryLocation::right:
				return clockwise ? corners[bottomRight] : corners[topRight];
			case BoundaryLocation::top:
				return clockwise ? corners[topRight]    : corners[topLeft];
		}
		throw std::runtime_error("Unexpected branch");
	};

	ArrayXi edges;
	auto addNewEdges = [&](){
		ArrayXi newEdges;
		if (
			(clockwise &&
			(locBeg == BoundaryLocation::left ||
				locBeg == BoundaryLocation::top) ) ||
			(!clockwise &&
			(locBeg == BoundaryLocation::bottom ||
				locBeg == BoundaryLocation::right) )
		){
			newEdges = boundarySegment( locBeg, vBeg, vEnd, meshSize );
		} else {
			newEdges = boundarySegmentReverse( locBeg, vBeg, vEnd, meshSize );
		}
		if ( edges.size() == 0 )
			edges = std::move(newEdges);
		else {
			ArrayXi temp ( edges.size() + newEdges.size() );
			temp << edges, newEdges;
			edges = std::move(temp);
		}
	};

	bool isSameSide {false};
	while ( !isSameSide ){
		isSameSide = sameSide();
		
		if ( !isSameSide ){
			vEnd = nextCorner();
		}
		addNewEdges();
		progressLoc();
		vBeg = vEnd;
		vEnd = verts[2];
	}

	return edges;
}

namespace vertsToPatches
{


void assertVerticesAtStructBoundary(
	const Array3i& vertices,
	const Ref<const Array2i>& meshSize,
	const Array4i& corners
){
	Index nx { meshSize.x() };

	/* check 0: vertices must be in mesh */
	if ( (vertices > corners[topRight]).any() || (vertices < 0).any() )
		throw InputError(
			"Boundary Patch: At least one vertex is not part of the mesh!"
		);
	
	/* check 1: vertices must be boundary vertices */
	for ( auto vertex : vertices ){
		Index rem { vertex % (nx+1) };
		if (
			vertex > corners[bottomRight] &&
			vertex < corners[topLeft] &&
			rem != 0 && rem != nx
		)
			throw InputError("Boundary Patch: Vertices not at boundary!");
	}

	/* check 2: vertex 2 must be next to vertex 1 */
	Index diff { vertices[1] - vertices[0] };
	if ( abs(diff) != 1 && abs(diff) != nx+1 )
		throw InputError(
			"Boundary Patch: Second vertex is not next to first vertex!"
		);
}

BoundaryLocation whereAre(
	const Ref<const ArrayXi>& vertices,
	const Ref<const Array2i>& meshSize,
	const Array4i& corners
){
	Index nx { meshSize.x() };

	using BL = BoundaryLocation;

	if ( (vertices <= corners[bottomRight]).all() )
		return BL::bottom;
	else if ( (vertices >= corners[topLeft]).all() )
		return BL::top;
	else {
		/* modulo as unary expression */
		ArrayXi remainders { vertices.unaryExpr(
			[&](Index coeff){ return coeff % (nx+1); }
		) };
		
		if ( (remainders == 0).all() )
			return BL::left;
		else if ( (remainders == nx).all() )
			return BL::right;
		else
			throw std::logic_error("Vertices not on the same side!");
	}
}

bool isClockwise ( const Array3i& vertices, BoundaryLocation loc ){
	using BL = BoundaryLocation;
	return ( vertices[1] > vertices[0] ) ?
		(loc == BL::left   || loc == BL::top  ) :
		(loc == BL::bottom || loc == BL::right);
}

std::optional<Corner> endCorner(
	Index endVertex,
	const Array4i& corners
){
	for ( Index i {0}; i<4; ++i ){
		if ( corners[i] == endVertex ){
			return static_cast<Corner>(i);
		}
	}
	return std::nullopt;
}

BoundaryLocation endLocation(
	Index endVertex,
	const Ref<const Array2i>& meshSize,
	const Array4i& corners,
	bool clockwise
){
	using BL = BoundaryLocation;

	std::optional<Corner> endIsCorner { endCorner(endVertex, corners) };

	if (endIsCorner){
		switch ( *endIsCorner ){
			case bottomLeft:
				return clockwise ? BL::bottom : BL::left;
			case bottomRight:
				return clockwise ? BL::right  : BL::bottom;
			case topLeft:
				return clockwise ? BL::left   : BL::top;
			case topRight:
				return clockwise ? BL::top    : BL::right;
			default:
				;
		}
	}
	return whereAre( Array1i{endVertex}, meshSize, corners );
}

BoundaryLocation nextAfter( BoundaryLocation loc, bool clockwise ){
	using BL = BoundaryLocation;
	switch (loc){
		case BL::bottom: return clockwise ? BL::left   : BL::right;
		case BL::top   : return clockwise ? BL::right  : BL::left;
		case BL::left  : return clockwise ? BL::top    : BL::bottom;
		case BL::right : return clockwise ? BL::bottom : BL::top;
		default: UNREACHABLE;
	}
}

Corner nextCorner( BoundaryLocation loc, bool clockwise ){
	using BL = BoundaryLocation;
	switch (loc){
		case BL::bottom: return clockwise ? bottomLeft  : bottomRight;
		case BL::top   : return clockwise ? topRight    : topLeft;
		case BL::left  : return clockwise ? topLeft     : bottomLeft;
		case BL::right : return clockwise ? bottomRight : topRight;
		default: UNREACHABLE;
	}
}

} // namespace vertsToPatches

std::vector<ArrayXi> verticesToStructPatches(
	Array3i vertices,
	const Ref<const Array2i>& meshSize
){
	using namespace vertsToPatches;

	Array4i corners { structMeshCorners(meshSize) };
	/* make sure input is valid */
	assertVerticesAtStructBoundary(vertices, meshSize, corners);

	/* find location of first edge and rotational direction */
	using BL = BoundaryLocation;
	BL locCur { whereAre(vertices.head<2>(), meshSize, corners) };
	bool clockwise ( isClockwise(vertices, locCur) );
	/* with that, we can find the location of the last vertex. Without the 
	 * rotational direction, corners are ambivalent */
	BL locEnd { endLocation( vertices[2], meshSize, corners, clockwise ) };

	std::vector<ArrayXi> patches;

	Index vBeg { vertices[0] }, vEnd {vertices[2] };

	auto addPatch = [&](){
		patches.emplace_back( boundarySegment(
			locCur, std::min(vBeg,vEnd), std::max(vEnd,vBeg), meshSize
		) );
	};

	/* if start and end are on the same side, but we have to go around to get
	 * to the end, then locCur == locEnd, but the loop mustn't be skipped */
	bool clockwiseAscendingSide { locCur == BL::left || locCur == BL::top };
	bool ascending { vertices[0] < vertices[2] };
	bool goAroundAnyway {
		locCur == locEnd && (
			vertices[0] == vertices[2] ||
			/* (A && B) || (!A && !B) being equal to A == B gives
			 * the expression below, after some trying it out on paper */
			clockwise == (clockwiseAscendingSide != ascending)
		)
	};

	while( locCur != locEnd || goAroundAnyway ){
		goAroundAnyway = false; // no infinite loop
		vEnd = corners[nextCorner(locCur, clockwise)];
		addPatch();	
		locCur = nextAfter(locCur, clockwise);
		vBeg = vEnd;
		vEnd = vertices[2];
	}
	addPatch();

	return patches;
}



template<typename MeshType>
Array3i boundarySegmentToVertices(
	const MeshBase<MeshType>& mesh,
	const ArrayXi& boundarySegment
){
	assert( boundarySegment.size() > 0 );
	/* if the boundary segment is very short, return early */
	if ( boundarySegment.size() == 1 ){
		Array2i edgeVertices { mesh.edge( boundarySegment[0] ).vertices() };
		return { edgeVertices[0], edgeVertices[1], edgeVertices[1] };
	}

	/* first entry holds the first, second the second and third the last vertex
	 * of that boundary segment */
	Array3i vertices;

	/* the second and second last vertices are the ones shared by the first and
	 * last two edges, respectively */
	vertices[1] = sharedVertex(
		mesh.edge( boundarySegment[0] ),
		mesh.edge( boundarySegment[1] )
	);

	auto lastTwoEdges { boundarySegment.tail(2) };

	Index secondLast = sharedVertex(
		mesh.edge( lastTwoEdges[0] ),
		mesh.edge( lastTwoEdges[1] )
	);

	/* the first and last vertex are the respective other (=not shared) */
	vertices[0] = otherVertex(
		mesh.edge( boundarySegment[0] ),
		vertices[1]
	);

	vertices[2] = otherVertex(
		mesh.edge( boundarySegment.tail(1)[0] ),
		secondLast
	);

	return vertices;
}


template<typename MeshType>
ArrayXi verticesToBoundarySegment(
	const MeshBase<MeshType>& mesh,
	const Array3i& verts,
	const Array3Xi& edgesAtVertex
){
	if constexpr (
		std::is_same_v<MeshType, Mesh>
	){
		assert(
			edgesAtVertex.size() != 0 &&
			"Edges at boundary vertices must be passed in for unstructured meshes!"
		);
		return verticesToBoundarySegment(
			mesh.derived().edges().vertices(), edgesAtVertex, verts
		);
	} else {
		return verticesToBoundarySegment(mesh.derived().size(), verts);
	}
}


template<typename MeshType>
void addBoundaryPatch(
	MeshBase<MeshType>& mesh,
	const std::string& patchName,
	const Array3i& vertices,
	const Array3Xi& edgesAtVertex
){
	/* only boundary patches that do not cross corners are allowed */
	if constexpr ( std::is_base_of_v<StructMeshBase<MeshType>, MeshType> ){

		std::vector<ArrayXi> patchEdges {
			verticesToStructPatches( vertices, mesh.derived().size() )
		};
		if (patchEdges.size() == 1){
			mesh.addBoundaryPatch(
				patchName, std::move(patchEdges[0])
			);
		} else {
			for (size_t i{0}; i<patchEdges.size(); ++i){
				mesh.addBoundaryPatch(
					patchName + "_" + std::to_string(i),
					std::move(patchEdges[i])
				);
			}
		}
	} else {
		mesh.addBoundaryPatch(
			patchName,
			verticesToBoundarySegment( mesh, vertices, edgesAtVertex )
		);
	}
}


#define INSTANTIATE_FOR_ALL_MESHTYPES(MeshType) \
template Array3i boundarySegmentToVertices( \
	const MeshBase<MeshType>&, const ArrayXi& ); \
template ArrayXi verticesToBoundarySegment( \
	const MeshBase<MeshType>&, const Array3i& verts, const Array3Xi& ); \
template void addBoundaryPatch( \
	MeshBase<MeshType>&, const std::string&, const Array3i&, const Array3Xi& );

#include "meshInstantiator.hpp"

} // end namespace hms
